from pc import data_out
import cv2 as cv
import numpy as np
from ultralytics import YOLO
import time
print("start")
model=YOLO("gta.pt")
cap=cv.VideoCapture(0)


def cvon():
    ret, frame=cap.read()
    results=model(frame)
    for result in results:
        boxes=result.boxes
        if boxes is not None:
                for box in boxes:
                    # 获取边界框坐标
                    x1, y1, x2, y2 = map(int, box.xyxy[0])
                    # 获取置信度
                    conf = float(box.conf[0])
                    # 获取类别
                    cls = int(box.cls[0])
                    # 获取类别名称
                    cls_name = model.names[cls]
                    
                    # 绘制边界框和标签
                    if conf > 0.5:  # 只显示置信度大于0.5的检测结果
                        # 绘制边界框
                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                        # 创建标签文本
                        label = f"{cls_name}: {conf:.2f}"
                        # 计算标签文本大小
                        (label_width, label_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
                        # 绘制标签背景
                        cv2.rectangle(frame, (x1, y1 - label_height - baseline), 
                                    (x1 + label_width, y1), (0, 255, 0), -1)
                        # 绘制标签文本
                        cv2.putText(frame, label, (x1, y1 - baseline), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
        cv2.imshow("frame",frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        
if cap.isOpened():
    print("open")
    cap.read()
    ret, frame=cap.read()
    if not ret:
        print("error")
    else:
        while True:
            cvon()
else:
    print("error")




        

